import psycopg2
import psycopg2.extras
import re

class KingbaseDB:
    def __init__(self, config):
        self.config = config
        self.connection = None
        self.logger = None  # 初始化logger属性
    
    def set_logger(self, logger):
        """
        设置日志记录器
        """
        self.logger = logger
    
    def connect(self):
        """
        连接到KingbaseES数据库
        """
        try:
            self.connection = psycopg2.connect(
                host=self.config['host'],
                port=self.config['port'],
                user=self.config['user'],
                password=self.config['password'],
                dbname=self.config['database']
            )
            return True
        except Exception as e:
            raise Exception(f"KingbaseES数据库连接失败: {e}")
    
    def disconnect(self):
        """
        断开数据库连接
        """
        if self.connection and not self.connection.closed:
            self.connection.close()
    
    def execute_query(self, query, params=None):
        """
        执行查询语句并返回结果
        """
        try:
            with self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                cursor.execute(query, params)
                result = cursor.fetchall()
                return result
        except Exception as e:
            raise Exception(f"执行查询失败: {query}\n{params}\n{e}")
    
    def execute_update(self, query, params=None):
        """
        执行更新语句（INSERT, UPDATE, DELETE等）
        """
        try:
            with self.connection.cursor() as cursor:
                affected_rows = cursor.execute(query, params)
                self.connection.commit()
                return affected_rows
        except Exception as e:
            self.connection.rollback()
            raise Exception(f"执行更新失败: {query}\n{params}\n{e}")
    
    def execute_batch(self, query, params_list):
        """
        批量执行SQL语句
        """
        try:
            with self.connection.cursor() as cursor:
                # 检查是否是INSERT语句，如果是，则尝试添加ON CONFLICT DO NOTHING来处理主键冲突
                if query.strip().upper().startswith('INSERT'):
                    # 使用正则表达式查找VALUES关键字，避免误识别列名中的"VALUES"
                    # 匹配独立的VALUES关键字，前后有空格或括号
                    values_match = re.search(r'\bVALUES\b', query, re.IGNORECASE)
                    if values_match:
                        values_pos = values_match.start()
                        # 构建支持ON CONFLICT DO NOTHING的INSERT语句
                        insert_with_conflict = query[:values_pos] + ' VALUES' + query[values_pos+6:] + ' ON CONFLICT DO NOTHING'
                        psycopg2.extras.execute_batch(cursor, insert_with_conflict, params_list)
                    else:
                        psycopg2.extras.execute_batch(cursor, query, params_list)
                else:
                    psycopg2.extras.execute_batch(cursor, query, params_list)
                self.connection.commit()
        except Exception as e:
            # 如果是主键冲突错误，尝试单独处理每条记录
            if 'duplicate key value' in str(e) or 'conflict' in str(e).lower():
                if self.logger:
                    self.logger.warning(f"批量插入遇到主键冲突，尝试单独处理每条记录: {str(e)}")
                self.connection.rollback()
                
                # 单独处理每条记录
                conflict_count = 0
                success_count = 0
                for params in params_list:
                    try:
                        with self.connection.cursor() as single_cursor:
                            # 尝试添加ON CONFLICT DO NOTHING
                            if query.strip().upper().startswith('INSERT'):
                                # 使用正则表达式查找VALUES关键字
                                values_match = re.search(r'\bVALUES\b', query, re.IGNORECASE)
                                if values_match:
                                    values_pos = values_match.start()
                                    insert_with_conflict = query[:values_pos] + ' VALUES' + query[values_pos+6:] + ' ON CONFLICT DO NOTHING'
                                    single_cursor.execute(insert_with_conflict, params)
                                else:
                                    single_cursor.execute(query, params)
                            else:
                                single_cursor.execute(query, params)
                            success_count += 1
                    except Exception as single_e:
                        # 如果是主键冲突，跳过该记录
                        if 'duplicate key value' in str(single_e) or 'conflict' in str(single_e).lower():
                            conflict_count += 1
                            continue
                        # 如果是其他错误，重新抛出
                        self.connection.rollback()
                        raise Exception(f"单独执行SQL失败: {query}\n{params}\n{single_e}")
                
                self.connection.commit()
                if self.logger:
                    self.logger.info(f"处理完成: {success_count}条记录成功，{conflict_count}条记录因主键冲突被跳过")
            else:
                self.connection.rollback()
                raise Exception(f"批量执行失败: {query}\n{e}")
    
    def create_schema(self, schema_name):
        """
        创建模式
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        query = f"CREATE SCHEMA IF NOT EXISTS {schema_name}"
        try:
            self.execute_update(query)
            return True
        except Exception as e:
            raise Exception(f"创建模式失败: {schema_name}\n{e}")
    
    def drop_schema(self, schema_name, cascade=False):
        """
        删除模式
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        cascade_sql = " CASCADE" if cascade else ""
        query = f"DROP SCHEMA IF EXISTS {schema_name}{cascade_sql}"
        try:
            self.execute_update(query)
            return True
        except Exception as e:
            raise Exception(f"删除模式失败: {schema_name}\n{e}")
    
    def table_exists(self, schema_name, table_name):
        """
        检查表是否存在
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        # PostgreSQL/KingbaseES在information_schema中存储的是小写表名
        # 所以我们需要将传入的表名也转换为小写进行比较
        query = """
        SELECT EXISTS (
            SELECT 1 
            FROM information_schema.tables 
            WHERE table_schema = %s 
            AND table_name = lower(%s)
        )
        """
        result = self.execute_query(query, (schema_name, table_name))
        if result and len(result) > 0:
            return result[0]['exists']
        return False
    
    def drop_table(self, schema_name, table_name):
        """
        删除表
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        # 表名用双引号括起来以保持大小写
        query = f"DROP TABLE IF EXISTS {schema_name}.\"{table_name}\""
        try:
            self.execute_update(query)
            return True
        except Exception as e:
            raise Exception(f"删除表失败: {schema_name}.{table_name}\n{e}")
    
    def create_table(self, schema_name, create_table_sql):
        """
        创建表
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        try:
            # 先检查并创建模式
            self.create_schema(schema_name)
            # 执行创建表SQL
            self.execute_update(create_table_sql)
            return True
        except Exception as e:
            raise Exception(f"创建表失败: {schema_name}\n{create_table_sql}\n{e}")
    
    def get_table_count(self, schema_name, table_name):
        """
        获取表的记录数
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        # 表名用双引号括起来以保持大小写
        query = f'SELECT COUNT(*) AS count FROM {schema_name}."{table_name}"'
        result = self.execute_query(query)
        if result and len(result) > 0:
            return result[0]['count']
        return 0
    
    def get_all_tables_in_schema(self, schema_name):
        """
        获取模式下的所有表名
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        query = """
        SELECT 
            table_name 
        FROM 
            information_schema.tables 
        WHERE 
            table_schema = %s 
            AND table_type = 'BASE TABLE'
        """
        result = self.execute_query(query, (schema_name,))
        return [row['table_name'] for row in result]
    
    def create_index(self, schema_name, index_name, table_name, columns, is_unique=False):
        """
        创建索引
        """
        schema_name = schema_name.lower()  # 确保模式名为小写
        try:
            # 构建索引类型关键字
            index_type = 'UNIQUE' if is_unique else ''
            
            # 构建索引列定义
            columns_str = ', '.join(columns)
            
            # 构建创建索引的SQL
            create_index_sql = f'CREATE {index_type} INDEX IF NOT EXISTS {index_name} ON {schema_name}."{table_name}" ({columns_str});'
            
            # 执行SQL
            self.execute_update(create_index_sql)
            return True
        except Exception as e:
            raise Exception(f"创建索引失败: {schema_name}.{table_name}.{index_name}\n{e}")